#include <opencv2/opencv.hpp>
#include <vector>
#include <iostream>
#include <string>

using std::vector;
using std::string;
using std::cout;
using std::endl;


int main(){
    int BOARD_SIZE[2] = {6,9};

    string img_path = "../image1/*.jpg";
    vector<cv::String> imgs_path;
    cv::glob(img_path, imgs_path);

    vector<cv::Point3f> pts_w; //棋盘点的世界坐标
    for(int i=0;i<BOARD_SIZE[1];i++)
        for(int j=0;j<BOARD_SIZE[0];j++)
            pts_w.push_back(cv::Point3f(i,j,0));
    
    vector<vector<cv::Point2f>> imgs_corners; //所有图片的棋盘点
    vector<cv::Point2f> pts_img; //单张图片的棋盘点
    vector<vector<cv::Point3f>> pts_w_all; //每张图片棋盘点对应的世界坐标

    cv::Mat img;
    cv::Mat img_gray;

    for(auto path: imgs_path){
        img=cv::imread(path);
        cv::cvtColor(img, img_gray, cv::COLOR_BGR2GRAY);
        bool found = cv::findChessboardCorners(img_gray, cv::Size(BOARD_SIZE[0],BOARD_SIZE[1]),pts_img);

        if(found){
            cout<<"find chessboard corners: "<<pts_img.size()<<endl;

            //下面提取亚像素角点
            cv::TermCriteria criteria(CV_TERMCRIT_EPS | CV_TERMCRIT_ITER, 30, 0.001);
            cv::cornerSubPix(img_gray, pts_img, cv::Size(11,11), cv::Size(-1,-1), criteria);

            cv::drawChessboardCorners(img, cv::Size(BOARD_SIZE[0],BOARD_SIZE[1]), pts_img, found);

            imgs_corners.emplace_back(pts_img);
            pts_w_all.emplace_back(pts_w);
        }
        cv::imshow("found corners", img);
        cv::waitKey(0);

        pts_img.clear();
    }

    //接下来标定畸变
    cv::Mat cameraMatrix, distCoeffs, rvecs, tvecs;
    cv::calibrateCamera(pts_w_all, imgs_corners, img_gray.size(), cameraMatrix, distCoeffs, rvecs, tvecs, cv::CALIB_THIN_PRISM_MODEL);
    //cv::calibrateCamera(pts_w_all, imgs_corners, img_gray.size(), cameraMatrix, distCoeffs, rvecs, tvecs);

    cout<<"cameraMatrix"<<endl<<cameraMatrix<<endl;
    cout<<"distCoeffs"<<endl<<distCoeffs<<endl;
    cout<<"rvecs"<<endl<<rvecs<<endl;
    cout<<"tvecs"<<endl<<tvecs<<endl;

    //尝试校准畸变图像
    for(auto path: imgs_path){
        img=cv::imread(path);
        cv::Mat img_undt;
        cv::undistort(img, img_undt, cameraMatrix, distCoeffs);
        cv::Mat compare;
        cv::vconcat(img,img_undt, compare);

        cv::imshow("campare", compare);
        cv::waitKey(0);
    }
}